package com.sky.utils;

import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.bouncycastle.cert.X509v3CertificateBuilder;
import org.bouncycastle.cert.jcajce.JcaX509CertificateConverter;
import org.bouncycastle.cert.jcajce.JcaX509v3CertificateBuilder;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.operator.ContentSigner;
import org.bouncycastle.operator.jcajce.JcaContentSignerBuilder;
import org.springframework.stereotype.Component;

import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.math.BigInteger;
import java.security.*;
import java.security.cert.Certificate;
import org.bouncycastle.asn1.x500.X500Name;

import javax.annotation.PostConstruct;
import java.security.cert.X509Certificate;
import java.util.Base64;
import java.util.Date;

/**
 * KeyStore 工具类，管理用户的密钥对和公钥。
 */
@Slf4j
@Component
@Data
public class KeyStoreUtil {

    private static final String KEYSTORE_TYPE = "PKCS12"; // 可选择 "JKS" 或其他类型
    private static final String KEYSTORE_PATH = "user-keystore.p12";
    private static final char[] KEYSTORE_PASSWORD = "changeit".toCharArray();

    /**
     * 初始化 KeyStore，如果不存在则创建一个新的空 KeyStore
     */
    public void initKeyStore() throws Exception {
        KeyStore keyStore = KeyStore.getInstance(KEYSTORE_TYPE);
        try (FileInputStream fis = new FileInputStream(KEYSTORE_PATH)) {
            keyStore.load(fis, KEYSTORE_PASSWORD); // 加载现有 KeyStore
        } catch (Exception e) {
            // 如果文件不存在，则创建一个空的 KeyStore
            keyStore.load(null, null);
            try (FileOutputStream fos = new FileOutputStream(KEYSTORE_PATH)) {
                keyStore.store(fos, KEYSTORE_PASSWORD);
            }
        }
    }

    @PostConstruct // Spring 注解，表示在 Bean 初始化后调用
    public void initialize() {
        try {
            initKeyStore();
            log.info("KeyStore 初始化成功!!!");
        } catch (Exception e) {
            throw new RuntimeException("KeyStore 初始化失败", e);
        }
    }

    /**
     * 为用户生成 RSA 密钥对并保存到 KeyStore
     *
     * @param userId 用户唯一标识
     */
    public void generateAndSaveKeyPair(String userId) throws Exception {
        KeyStore keyStore = KeyStore.getInstance(KEYSTORE_TYPE);
        try (FileInputStream fis = new FileInputStream(KEYSTORE_PATH)) {
            keyStore.load(fis, KEYSTORE_PASSWORD);
        }

        // 生成密钥对
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance("RSA");
        keyPairGenerator.initialize(2048);
        KeyPair keyPair = keyPairGenerator.generateKeyPair();

        // 生成自签名证书
        String dn = "CN=" + userId + ", O=MyCompany, C=US"; // 证书的主题名
        X509Certificate certificate = generateSelfSignedCertificate(keyPair, dn, 365);

        // 存储私钥和证书
        keyStore.setKeyEntry(userId, keyPair.getPrivate(), KEYSTORE_PASSWORD, new Certificate[]{certificate});

        // 保存 KeyStore
        try (FileOutputStream fos = new FileOutputStream(KEYSTORE_PATH)) {
            keyStore.store(fos, KEYSTORE_PASSWORD);
        }
    }

    /**
     * 加载用户的密钥对
     *
     * @param userId 用户唯一标识
     * @return 密钥对
     */
    public KeyPair loadKeyPair(String userId) throws Exception {
        KeyStore keyStore = KeyStore.getInstance(KEYSTORE_TYPE);
        try (FileInputStream fis = new FileInputStream(KEYSTORE_PATH)) {
            keyStore.load(fis, KEYSTORE_PASSWORD);
        }

        PrivateKey privateKey = (PrivateKey) keyStore.getKey(userId, KEYSTORE_PASSWORD);
        Certificate certificate = keyStore.getCertificate(userId);
        PublicKey publicKey = certificate.getPublicKey();

        return new KeyPair(publicKey, privateKey);
    }

    /**
     * 保存前端传回的公钥到 KeyStore
     *
     * @param userId           用户唯一标识
     * @param base64PublicKey  Base64 编码的公钥字符串
     */
    public void saveClientPublicKey(String userId, String base64PublicKey) throws Exception {
        KeyStore keyStore = KeyStore.getInstance(KEYSTORE_TYPE);
        try (FileInputStream fis = new FileInputStream(KEYSTORE_PATH)) {
            keyStore.load(fis, KEYSTORE_PASSWORD);
        }

        byte[] publicKeyBytes = Base64.getDecoder().decode(base64PublicKey);
        PublicKey publicKey = KeyFactory.getInstance("RSA")
                .generatePublic(new java.security.spec.X509EncodedKeySpec(publicKeyBytes));

        // 生成证书对象（仅用于演示，生产环境请使用 CA 签发证书）
        String dn = "CN=" + userId + ", O=Client, C=US";
        X509Certificate certificate = generateCertificateFromPublicKey(publicKey, dn);

        // 存储公钥为证书
        keyStore.setCertificateEntry(userId, certificate);

        // 保存 KeyStore
        try (FileOutputStream fos = new FileOutputStream(KEYSTORE_PATH)) {
            keyStore.store(fos, KEYSTORE_PASSWORD);
        }
    }

    /**
     * 加载前端保存的公钥
     *
     * @param userId 用户唯一标识
     * @return 公钥对象
     */
    public PublicKey loadClientPublicKey(String userId) throws Exception {
        KeyStore keyStore = KeyStore.getInstance(KEYSTORE_TYPE);
        try (FileInputStream fis = new FileInputStream(KEYSTORE_PATH)) {
            keyStore.load(fis, KEYSTORE_PASSWORD);
        }

        Certificate certificate = keyStore.getCertificate(userId);
        if (certificate != null) {
            return certificate.getPublicKey();
        }
        throw new Exception("公钥未找到: " + userId);
    }

    private X509Certificate generateSelfSignedCertificate(KeyPair keyPair, String dn, int days) throws Exception {
        Security.addProvider(new BouncyCastleProvider());

        // 设置证书信息
        X500Name issuerName = new X500Name(dn);
        X500Name subjectName = issuerName; // 自签名证书，颁发者和主题相同
        Date startDate = new Date();
        Date endDate = new Date(startDate.getTime() + days * 24L * 60 * 60 * 1000L);
        BigInteger serialNumber = new BigInteger(64, new SecureRandom());

        // 创建证书生成器
        X509v3CertificateBuilder certBuilder = new JcaX509v3CertificateBuilder(
                issuerName,
                serialNumber,
                startDate,
                endDate,
                subjectName,
                keyPair.getPublic()
        );

        // 签名证书
        ContentSigner signer = new JcaContentSignerBuilder("SHA256withRSA")
                .build(keyPair.getPrivate());

        // 转换为 X.509 证书
        return new JcaX509CertificateConverter()
                .setProvider("BC")
                .getCertificate(certBuilder.build(signer));
    }

    /**
     * 从公钥生成简单的证书（仅示例，推荐使用 BouncyCastle）
     */
    private X509Certificate generateCertificateFromPublicKey(PublicKey publicKey, String dn) throws Exception {
        KeyPair tempKeyPair = new KeyPair(publicKey, null); // 临时 KeyPair
        return generateSelfSignedCertificate(tempKeyPair, dn, 365);
    }
}
